"""
RunnableLambda和注解chain使用示例
"""

from operator import itemgetter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import chain, RunnableLambda

from models import get_ds_model_client


def length_fn(text):
    return len(text)

def _multi_length_fn(text1, text2):
    return len(text1) * len(text2)

@chain
def multi_length_chain_fn(_dict):
    return _multi_length_fn(_dict["text1"], _dict["text2"])

prompt = ChatPromptTemplate.from_template("{a} + {b}是什么?")
client = get_ds_model_client()

chain1 = prompt | client

chain2 = (
    {
        "a": (itemgetter("foo") | RunnableLambda(length_fn)),
        "b": ({"text1": itemgetter("foo"), "text2": itemgetter("bar")} | multi_length_chain_fn)
    } | chain1
)
print(chain2.invoke({"foo": "hello", "bar": "world"}))



